
#pragma once

#include "tbb/tick_count.h"
#include "tbb/task.h"
#include "tbb/task_scheduler_init.h"
#include "BinaryTree.h"
#include <vector>

using namespace tbb;

template <class T>
class tbbTreeSumTask: public tbb::task
{
    T *res;
    BinaryTree<T>* root;
public:
    tbbTreeSumTask(T* res, BinaryTree<T> *root)
	{
		this->res = res;
		this->root = root;
	}

	tbb::task* execute()
	{
		if( root->getChildCount() < 256 * 1024 )
			*res = root->getSum();
		else
		{
            T x, y;
            int count = 1; 
            tbb::task_list list;
			if(root->getLeft() != NULL)
			{
                ++count;
                list.push_back( *new( allocate_child() ) tbbTreeSumTask(&x, root->getLeft()) );
            }
			if(root->getRight() != NULL)
			{
                ++count;
                list.push_back( *new( allocate_child() ) tbbTreeSumTask(&y, root->getRight()) );
            }
            set_ref_count(count);
            spawn_and_wait_for_all(list);
			*res = root->getValue();
			if(root->getLeft() != NULL)
				*res += x;
			if(root->getRight() != NULL)
				*res += y;
        }
        return NULL;
    }
};

//Parallel version of tree summator (via tbb). Return result + elapsed time.
template <class T>
std::pair<T, double> tbbTreeSum(BinaryTree<T> *tree)
{
	T res;
    tbbTreeSumTask<T> &a = *new(tbb::task::allocate_root()) tbbTreeSumTask<T>(&res, tree);
	tick_count t_start = tick_count::now();
    tbb::task::spawn_root_and_wait(a);
	tick_count t_end = tick_count::now();
	return std::pair<T, double>(res, (t_end - t_start).seconds());
}
